// Copyright 2023 Innkeeper Lukexp(王晓平) &lt;37196496@qq.com>. All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file. The original repo for
// this file is https://gitee.com/jetxt/jinyu.

package worker

import (
	"sync"
)

const (
	qsize = 200
)

type TaskFunc func()

type task struct {
	fn TaskFunc
}

func (t *task) Do() {
	t.fn()
}

type worker struct {
	q       chan *task
	subsize int
	stopCh  chan struct{}
	wg      sync.WaitGroup
}

func NewWorker(size int) *worker {
	wp := &worker{
		q:       make(chan *task, qsize),
		subsize: size,
		stopCh:  make(chan struct{}),
		wg:      sync.WaitGroup{},
	}
	wp.start()
	return wp
}

func (wp *worker) Add(f TaskFunc) {
	wp.wg.Add(1)
	wp.q <- &task{
		fn: f,
	}
}

func (wp *worker) Wait() {
	wp.wg.Wait()
	wp.stop()
}

func (wp *worker) start() {
	for i := 0; i < wp.subsize; i++ {
		go func() {
			for {
				select {
				case t := <-wp.q:
					t.Do()
					wp.wg.Done()
				case <-wp.stopCh:
					return
				}
			}
		}()
	}
}

func (wp *worker) stop() {
	close(wp.stopCh)
}
